"""
边与节点特征点积操作

提供与DGL兼容的边与节点特征点积操作，使用PyG作为后端
"""

import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu
import sys
import os

# 添加路径
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(os.path.dirname(current_dir))
if parent_dir not in sys.path:
    sys.path.insert(0, parent_dir)

from ops.src.graph.graph import is_dgl_graph


def e_dot_v(graph, edge_features, node_features):
    """
    计算边特征与节点特征的点积
    
    参数:
        graph: Graph对象或DGL图
        edge_features: 边特征张量
        node_features: 节点特征张量
    
    返回:
        边与节点特征的点积结果
    """
    # 检查图是否为DGL图
    # print("===================e_dot_v===================")
    if is_dgl_graph(graph):
        # DGL图
        import dgl.function as fn
        graph = graph.local_var()
        graph.edata['e'] = edge_features
        graph.ndata['v'] = node_features
        graph.apply_edges(fn.e_dot_v('e', 'v', 'out'))
        return graph.edata['out']
    else:
        # PyG风格的Graph
        # 获取边的源节点和目标节点索引
        src, dst = graph.edge_index
        
        # 获取源节点的特征
        # 注意：DGL的e_dot_v是使用目标节点的特征，而不是源节点
        dst_features = node_features[dst]
        
        # 计算点积
        # 对齐DGL的实现，计算 sum(e_i * v_j, dim=-1)
        result = torch.sum(edge_features * dst_features, dim=-1, keepdim=True)
        
        return result 